# -*- coding: utf-8 -*-
"""
Project: P2PServer
File: Server.py
Author:  laoknas
Created date: 2025/02/08
Modified date: 2025/02/08
Comment: Main server UI for P2P network.
"""
try:
    import tkinter as tk
    from tkinter import ttk, scrolledtext
except ImportError as e:
    # print(f"缺少必要的库: {e}")
    pass
import socket
import threading
from cryptography.fernet import Fernet
import argparse
import logging
import os
from datetime import datetime
import base64
from pathlib import Path
from icon import img

class ServerUI:
    def __init__(self, host="0.0.0.0", port=12345, show_ui=True):
        self.host = host
        self.port = port
        self.show_ui = show_ui

        # 初始化日志功能
        self.init_logging()

        if self.show_ui:
            self.window = tk.Tk()
            self.window.title("P2P 服务器配置")
            # 设置icon
            try:
                # 使用base64编码，把导入DBSP.ICO文件进行编码，装入程序中。第一次编译是取消注释，当然也可以每次编译，都重新生产icon.py文件。
                # # 此处代码为把"favicon.ico"生成icon.py文件
                # open_icon = open("favicon.ico","rb")
                # b64str = base64.b64encode(open_icon.read())
                # open_icon.close()
                # write_data = "img = %s" % b64str
                # f = open("icon.py","w+")
                # f.write(write_data)
                # f.close()

                # # 此处代码为把"DBSP.ico"生成icon.py文件
                # from icon import img
                tmp = open("tmp.ico","wb+")
                tmp.write(base64.b64decode(img))
                tmp.close()
                self.window.iconbitmap("tmp.ico")
                Path("tmp.ico").unlink()
            except:
                pass
            self.create_widgets()

        self.server = None
        self.running = False
        self.clients = {}  # 存储客户端信息 {address: (socket, id, name)}
        self.key = Fernet.generate_key()  # 生成加密密钥
        self.cipher = Fernet(self.key)

    def init_logging(self):
        """初始化日志功能"""
        # 创建 logs 文件夹（如果不存在）
        if not os.path.exists("logs"):
            os.makedirs("logs")

        # 按日期生成日志文件名
        log_filename = datetime.now().strftime("logs/server_%Y-%m-%d.log")

        # 配置日志
        logging.basicConfig(
            level=logging.INFO,
            format="%(asctime)s - %(levelname)s - %(message)s",
            handlers=[
                logging.FileHandler(log_filename, encoding='utf-8'),  # 明确指定 UTF-8 编码
                logging.StreamHandler()  # 输出到控制台
            ]
        )

    def create_widgets(self):
        # 左侧配置区域
        left_frame = ttk.Frame(self.window)
        left_frame.grid(row=0, column=0, padx=10, pady=10, sticky="nsew")

        # 输入框：服务器地址
        lbl_host = ttk.Label(left_frame, text="服务器地址:")
        lbl_host.grid(row=0, column=0, padx=5, pady=5, sticky="w")
        self.entry_host = ttk.Entry(left_frame)
        self.entry_host.insert(0, self.host)
        self.entry_host.grid(row=0, column=1, padx=5, pady=5, sticky="ew")

        # 输入框：服务器端口
        lbl_port = ttk.Label(left_frame, text="服务器端口:")
        lbl_port.grid(row=1, column=0, padx=5, pady=5, sticky="w")
        self.entry_port = ttk.Entry(left_frame)
        self.entry_port.insert(0, str(self.port))
        self.entry_port.grid(row=1, column=1, padx=5, pady=5, sticky="ew")

        # 启动/停止按钮
        self.btn_toggle = ttk.Button(left_frame, text="启动服务器", command=self.toggle_server)
        self.btn_toggle.grid(row=2, column=0, columnspan=2, padx=5, pady=5, sticky="ew")

        # 关闭按钮
        self.btn_close = ttk.Button(left_frame, text="关闭", command=self.close_server)
        self.btn_close.grid(row=4, column=0, columnspan=2, padx=5, pady=5, sticky="ew")

        # 日志显示
        self.log_area = scrolledtext.ScrolledText(left_frame, height=15, width=50)
        self.log_area.grid(row=3, column=0, columnspan=2, padx=5, pady=5, sticky="nsew")

        # 右侧在线客户端列表
        right_frame = ttk.Frame(self.window)
        right_frame.grid(row=0, column=1, padx=10, pady=10, sticky="nsew")

        lbl_clients = ttk.Label(right_frame, text="在线客户端列表:")
        lbl_clients.grid(row=0, column=0, padx=5, pady=5, sticky="w")

        self.client_listbox = tk.Listbox(right_frame, width=30, height=20)
        self.client_listbox.grid(row=1, column=0, padx=5, pady=5, sticky="nsew")

        # 布局权重配置
        self.window.grid_columnconfigure(0, weight=1)
        self.window.grid_columnconfigure(1, weight=0)
        left_frame.grid_columnconfigure(1, weight=1)
        left_frame.grid_rowconfigure(3, weight=1)

    def log(self, message):
        if self.show_ui:
            self.log_area.insert(tk.END, message + "\n")
            self.log_area.see(tk.END)
        logging.info(message)  # 记录日志

    def start_server_thread(self):
        """启动服务器线程（修复非 UI 模式下的参数读取问题）"""
        if self.show_ui:
            # 从 UI 组件读取 host 和 port
            self.host = self.entry_host.get()
            self.port = int(self.entry_port.get())
        else:
            # 非 UI 模式下直接使用初始化参数
            self.host = self.host
            self.port = self.port

        try:
            self.server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.server.bind((self.host, self.port))
            self.server.listen(5)
            self.log(f"[*] 服务器启动，监听 {self.host}:{self.port}")
            self.running = True

            while self.running:
                try:
                    client_socket, client_address = self.server.accept()
                    self.log(f"[*] 新连接: {client_address}")
                    # 发送加密密钥给客户端
                    client_socket.send(self.key)
                    # 接收客户端ID和名称
                    client_info = client_socket.recv(1024).decode()
                    client_id, client_name = client_info.split(":")
                    self.clients[client_address] = (client_socket, client_id, client_name)
                    self.update_client_list()
                    # 广播更新后的在线客户端列表
                    self.broadcast_client_list()
                    client_handler = threading.Thread(target=self.handle_client, args=(client_socket, client_address))
                    client_handler.start()
                except Exception as e:
                    if self.running:
                        self.log(f"[!] 连接错误: {str(e)}")
                    break
        except Exception as e:
            self.log(f"[!] 服务器启动失败: {str(e)}")
            self.running = False

    def handle_client(self, client_socket, client_address):
        try:
            while True:
                data = client_socket.recv(1024)
                if not data:
                    break
                decrypted_data = self.cipher.decrypt(data).decode()
                # 新增消息类型判断
                if decrypted_data.startswith("PRIVATE:"):
                    # 格式: PRIVATE:target_id:message
                    _, target_id, message = decrypted_data.split(":", 2)
                    for addr, (sock, client_id, client_name) in self.clients.items():
                        if client_id == target_id:
                            sock.send(data)
                            break
                else:
                    # 广播给其他客户端，附加发送者名称
                    self.log(f"[*] 接收到来自 {self.clients[client_address][2]} 的消息: {decrypted_data}")
                    sender_name = self.clients[client_address][2]  # 获取发送者名称
                    broadcast_message = f"BROADCAST:{sender_name}:{decrypted_data}"
                    encrypted_broadcast = self.cipher.encrypt(broadcast_message.encode())
                    for addr, (sock, client_id, client_name) in self.clients.items():
                        if addr != client_address:
                            sock.send(encrypted_broadcast)
        except Exception as e:
            self.log(f"[!] 客户端 {self.clients[client_address][2]} 错误: {str(e)}")
        finally:
            client_socket.close()
            del self.clients[client_address]
            self.update_client_list()
            self.broadcast_client_list()

    def update_client_list(self):
        if self.show_ui:
            self.client_listbox.delete(0, tk.END)
            for client_address, (_, client_id, client_name) in self.clients.items():
                self.client_listbox.insert(tk.END, f"{client_name} ({client_id})")

    def broadcast_client_list(self):
        """广播在线客户端列表（修复数据格式）"""
        clients_data = ",".join([f"{client_id}:{client_name}" for _, client_id, client_name in self.clients.values()])
        broadcast_message = f"CLIENTS:{clients_data}"
        for client_socket, _, _ in self.clients.values():
            try:
                client_socket.send(self.cipher.encrypt(broadcast_message.encode()))
            except:
                pass

    def toggle_server(self):
        """切换服务器状态（修复线程启动逻辑）"""
        if not self.running:
            # 启动服务器线程（非守护线程，避免主线程退出后终止）
            server_thread = threading.Thread(target=self.start_server_thread)
            server_thread.start()
            if self.show_ui:
                self.btn_toggle.config(text="停止服务器")
        else:
            self.running = False
            if self.server:
                self.server.close()
            if self.show_ui:
                self.btn_toggle.config(text="启动服务器")
            self.log("[*] 服务器已停止")

    def close_server(self):
        """关闭服务器并退出程序"""
        self.running = False
        if self.server:
            self.server.close()
        if self.show_ui:
            self.window.destroy()  # 关闭 UI 窗口
        self.log("[*] 服务器已关闭")

    def run(self):
        if self.show_ui:
            self.window.mainloop()  # UI 模式进入主事件循环
        else:
            self.toggle_server()  # 非 UI 模式直接启动服务器
            try:
                while self.running:
                    # 非 UI 模式下保持主线程阻塞
                    threading.Event().wait(1)
            except KeyboardInterrupt:
                self.log("[*] 正在关闭服务器...")
                self.running = False
                if self.server:
                    self.server.close()


def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description="P2P 服务器")
    parser.add_argument("--host", type=str, default="0.0.0.0", help="服务器监听地址")
    parser.add_argument("--port", type=int, default=12345, help="服务器监听端口")
    parser.add_argument("--no-ui", action="store_true", help="不显示 UI 界面")
    return parser.parse_args()

def main():
    args = parse_args()
    server_ui = ServerUI(host=args.host, port=args.port, show_ui=not args.no_ui)
    server_ui.run()
    
if __name__ == "__main__":
    args = parse_args()
    server_ui = ServerUI(host=args.host, port=args.port, show_ui=not args.no_ui)
    server_ui.run()